# -*- codeing = utf-8 -*-
# @Time :2021/3/12 20:49
# @Author : 兢兢业业的小码农
# @File : word2vec_skipgram_01.py
# @Software : PyCharm
import warnings

# warnings.filterwarnings('ignore')

import os
import numpy as np
import tensorflow as tf
from python_ai.common.xcommon import *

# os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
# tf.compat.v1.logging.set_verbosity(40)

CORPUS = "natural language processing and machine learning is fun and exciting"
CORPUS = "I want to drink a glass of apple juice"

# 最大窗口数
window_size = 2
# 词嵌入矩阵的列数
embedding_dim = 5


# 数据处理
def dataSets(corpus_raw):
    # 去重后的字典
    words = []
    # 按空格切分后的句子
    sentences = []
    for word in corpus_raw.split():
        words.append(word)
        sentences.append(word)
    words = set(words)

    # 定义两个空字典
    word2idx = {}
    idx2word = {}

    for i, word in enumerate(words):
        # 存放单词对应的下标
        word2idx[word] = i
        # 存放下标对应的单词
        idx2word[i] = word

    train_data = []

    # 循环遍历每一个单词(不去重)
    for word_index, word in enumerate(sentences):
        # 移动窗口
        for nb_word in sentences[max(word_index - window_size, 0):min(word_index + window_size, len(sentences)) + 1]:
            # 判断窗口内的单词是否是目标词
            if nb_word != word:
                # 将目标词的下标和预测词的下标保存到一个列表
                train_data.append([word2idx[word], word2idx[nb_word]])
    return np.array(train_data), word2idx, idx2word


def word2vecMpodel(x_data, vocab_size):
    # 定义词嵌入矩阵
    W1 = tf.Variable(tf.random_normal([vocab_size, embedding_dim]))
    # 对数据进行操作
    hidden = tf.matmul(x_data, W1)
    # 定义权重
    W2 = tf.Variable(tf.random_normal([embedding_dim, vocab_size]))
    # 定义偏执
    b2 = tf.Variable(tf.random_normal([vocab_size]))
    # 预测结果
    logits = tf.matmul(hidden, W2) + b2
    return logits, W1


def mainProgram():
    train_data, word2idx, idx2word = dataSets(CORPUS)
    # print(train_data)

    # train_data 第一列是目标词 第二例是预测词
    x_train = np.reshape(train_data[:,0], [-1, 1])
    y_train = np.reshape(train_data[:,1], [-1, 1])

    # 是字典的长度 也是 键的长度
    vocab_size = len(word2idx.keys())
    print(vocab_size)
    # 定义占位符
    x = tf.placeholder(tf.int32, shape=[None, 1])
    y = tf.placeholder(tf.int32, shape=[None, 1])
    x_one_hot = tf.one_hot(x, vocab_size)
    y_one_hot = tf.one_hot(y, vocab_size)
    print(y_one_hot)
    print(x_one_hot)

    logits, W1 = word2vecMpodel(x_one_hot, vocab_size)

    loss = tf.reduce_mean(
        tf.nn.softmax_cross_entropy_with_logits(labels=y_one_hot, logits=logits)

    )

    optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.1).minimize(loss)

    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        for i in range(500):
            loss_Val, _ = sess.run([loss, optimizer], feed_dict={x: x_train, y: y_train})
            if i % 20 == 0:
                print(i + 1, 'loss', loss_Val)

        # 词嵌入矩阵
        embeddingMat = sess.run(W1)
        # 打印词嵌入矩阵
        sep('打印词嵌入矩阵')
        print('W1', embeddingMat)
        # 输出要打印的词 和 对应的词嵌入矩阵
        sep('输出要打印的词 和 对应的词嵌入矩阵')
        word = 'and'
        word = 'of'
        print(word, word2idx[word], embeddingMat[word2idx[word]])

        v_w1 = embeddingMat[word2idx[word]]

        # Check the logits of "want" ("to", "I")
        sep('Check the logits of "want" ("to", "I")')
        print(word2idx)
        print(idx2word)
        logits_v = sess.run(logits, feed_dict={x: [
            [word2idx['want']]
        ]})
        print(logits_v)

        # 选择排名前三位的
        sep('选择排名前三位的')
        top_n = 3
        word_sim = {}
        # 循环做比较
        for i in range(vocab_size):
            v_w2 = embeddingMat[i]
            theta_sum = np.dot(v_w1, v_w2)
            theta_dan = np.linalg.norm(v_w1) * np.linalg.norm(v_w2)
            # 余弦相似度
            theta = theta_sum / theta_dan
            word = idx2word[i]
            word_sim[word] = theta
        # 排序
        words_sorted = sorted(word_sim.items(), key=lambda kv: kv[1], reverse=True)
        print(words_sorted)
        for word, sim in words_sorted[:top_n]:
            print(word, sim)


if __name__ == '__main__':
    mainProgram()
